"""
# @Time : 2020/6/8 0008 23:10 
# @Author : xx
# @File : decode.py 
# @Software: PyCharm
"""
from concurrent.futures import ThreadPoolExecutor


def load_dict(dict_file_path, encoding="utf-8", blank_first=False):
    with open(dict_file_path, encoding=encoding, mode='r') as f:
        chars = list(map(lambda char: char.strip('\r\n'), f.readlines()))

    if blank_first:
        chars = chars[1:] + ['blank']

    dict = {i: v for i, v in enumerate(chars)}

    return dict


def decode_single_line(pred_text, nclass, id_to_char):
    char_list = []

    # pred_text = pred_text[np.where(pred_text != nclass - 1)[0]]

    for i in range(len(pred_text)):
        if pred_text[i] != nclass - 1 and (
                (pred_text[i] != pred_text[i - 1]) or (i > 1 and pred_text[i] == pred_text[i - 2])):
            char_list.append(id_to_char[pred_text[i]])
    return u''.join(char_list)


def split_single_line(pred, pred_text, nclass, id_to_char):
    split_index = []
    char_prob = []
    char_list = []
    # pred_text = pred_text[np.where(pred_text != nclass - 1)[0]]

    for i in range(len(pred_text)):
        if pred_text[i] != nclass - 1 and (
                (pred_text[i] != pred_text[i - 1]) or (i > 1 and pred_text[i] == pred_text[i - 2])):
            # split_index.append(i * w * scale / len(pred_text))
            split_index.append(i)
            char_prob.append(pred[i][pred_text[i]])
            char_list.append(id_to_char[pred_text[i]])
    return split_index, char_prob, char_list


def decode_split(preds, nclass, id_to_char):
    texts = []
    splits = []
    probs = []
    # start = time.time()

    pred_texts = preds.argmax(axis=2)

    with ThreadPoolExecutor() as executor:
        for split, prob, text in executor.map(
                lambda pred, pred_text:
                split_single_line(pred, pred_text, nclass, id_to_char),
                preds, pred_texts):
            texts.append(text)
            splits.append(split)
            probs.append(prob)
    # print("decode cost: ", time.time() - start)
    return texts, splits, probs


def decode(pred, dict_file_path):
    lines = []
    # print(type(pred), pred)
    id_to_char = load_dict(dict_file_path=dict_file_path, blank_first=False)
    nclass = id_to_char.keys().__len__()
    if isinstance(pred, list):
        import numpy as np
        pred = np.array(pred)
        # print(pred.shape)
    pred_texts = pred.argmax(axis=2)

    with ThreadPoolExecutor() as executor:
        for line in executor.map(lambda pred_text: decode_single_line(pred_text, nclass, id_to_char), pred_texts):
            lines.append(line)

    return lines